'''
- 图像切分
'''
import torch
import GLOB

def _image_segmentation(image,dtype=torch.float32):    # mabe convolution could help for smaller image
    # segment a single image with one channel
    picture_size=GLOB.PICTURE_SIZE
    segment_size=GLOB.SEGMENT_SIZE
    assert picture_size%segment_size==0    # picture should be cut inot segment evenly
    segment_num=int(picture_size/segment_size)
    segmentation=[]

    # implementation of linear position encoding, save a place
    position_encoding=[]
    for i in range(segment_num**2):
        position_encoding.append(torch.tensor([-1]))
    for row in range(segment_num):
        for col in  range(segment_num):
            segmentation.append(torch.cat((image[row*segment_size:(row+1)*segment_size,col*segment_size:(col+1)*segment_size].contiguous().view(-1),position_encoding[row*segment_num+col]),0))
            #print(segmentation[-1].size())
    tensor=torch.stack(segmentation)
    #print(tensor.size())
    return tensor
    

def batch_image_segmentation(batch_image, dtype=torch.float32):
    '''
    cut the image into segments and flatten each segments into a vector
    input: image size：(batch_size, num_channels, size, size)
    output: segmented images size:(batch_size, num_segments, len_segment)
    '''

    segment_size=GLOB.SEGMENT_SIZE
    # cut into segments
    segmented=[]
    for channel_image in batch_image:
        temp=[]
        for image in channel_image:
            temp.append(_image_segmentation(image=image, dtype=dtype))
        segmented.append(torch.stack(temp))
    tensor=torch.stack(segmented)
    # flatten color channels
    tensor=tensor.view(tensor.size()[0], -1, tensor.size(-1))
    
    #implement classify axis
    temp=[]
    for image in tensor:
        to_fix=torch.cat((torch.randn(size=[1,image.size()[-1]],dtype=torch.float32), image),dim=0)
        temp.append(to_fix)
    tensor=torch.stack(temp)    
    
    # position encoding
    for i in range(len(tensor)):
        for j in range(len(tensor[i])):
            tensor[i][j][-1]=float(2*(float(j)/len(tensor[i])))-1
    return tensor


